#include "control/move_control/lqr.h"
#include "control/control_params.h"

using namespace control;

void LQRPathFollowing::Init(ros::NodeHandle n)
{
    is_in_dec_stage_ = false;
    running_type_ = RunningType::HEADFOLLOWING;
    stop_distance_threshold_.SetDownThreshold(g_pose_control_accuracy);
    stop_distance_threshold_.SetUpThreshold(g_pose_control_accuracy * 3.0f);

    SetLQRQ();

    ptr_speed_smooth_ = std::make_shared<common_lib::SpeedSmoothWithCloseConstAcc>(0.05);
    ptr_anglespeed_smooth_ = std::make_shared<common_lib::SpeedSmoothWithCloseConstAcc>(0.05);
}

void LQRPathFollowing::SetLQRQ()
{
    LQR_Q_load_.setZero();
    LQR_Q_unload_.setZero();

    LQR_Q_load_(0, 0) = g_lat_error_weight_load;
    LQR_Q_load_(1, 1) = g_lat_error_rate_weight_load;
    LQR_Q_load_(2, 2) = g_heading_error_weight_load;
    LQR_Q_load_(3, 3) = g_heading_error_rate_weight_load;

    LQR_Q_unload_(0, 0) = g_lat_error_weight;
    LQR_Q_unload_(1, 1) = g_lat_error_rate_weight;
    LQR_Q_unload_(2, 2) = g_heading_error_weight;
    LQR_Q_unload_(3, 3) = g_heading_error_rate_weight;
    LINFO("LQR SetLQRQ() --> {:0.6f} {:0.6f} {:0.6f} {:0.6f}", g_lat_error_weight_load, g_lat_error_rate_weight_load, g_heading_error_weight_load, g_heading_error_rate_weight_load);
}

void LQRPathFollowing::VehicleController(amr_trajectory_msgs::amr_trajectory_msg &trajectory, AMRState &amr_state, float &speed, float &angle_speed, float &speed_toward, float &max_speed_limit, float &max_angle_speed_limit)
{
    static int fsm_switch_cout = 0;
    int fsm_switch_count_limit = 2;
    int fsm_head_angle_limit = 1;

    if (trajectory.points.empty())
    {
        VehicleStop();
        return;
    }
    if (!amr_state.actuator.is_auto_mode)
    {
        running_type_ = RunningType::HEADFOLLOWING;
    }

    if (running_type_ == RunningType::HEADFOLLOWING)
    {
        if (IsFollowingToRotating(trajectory, amr_state))
        {
            fsm_switch_cout++;
            if (fsm_switch_cout > fsm_switch_count_limit)
            {
                running_type_ = RunningType::ROTATING;
            }
            speed = 0;
            angle_speed = 0;
            ptr_speed_smooth_->Reset();
        }
        else
        {
            HeadFollowingController(trajectory, amr_state, speed, angle_speed, max_speed_limit, max_angle_speed_limit);
            fsm_switch_cout--;
        }
    }
    else
    {
        if (IsRotatingToFollowing(trajectory, amr_state))
        {
            fsm_switch_cout++;
            if (fsm_switch_cout > fsm_switch_count_limit)
            {
                running_type_ = RunningType::HEADFOLLOWING;
            }
            speed = 0;
            angle_speed = 0;
            ptr_speed_smooth_->Reset();
        }
        else
        {
            RotatingController(trajectory, amr_state, speed, angle_speed, max_angle_speed_limit);
            fsm_switch_cout--;
        }
    }

    if (fsm_switch_cout > fsm_head_angle_limit)
    {
        angle_speed = 0;
        speed = 0;
        ptr_speed_smooth_->Reset();
    }

    if (fsm_switch_cout < 0)
    {
        fsm_switch_cout = 0;
    }

    amr_state.running_type = running_type_;

    // ROS_INFO("VehicleController() --> runningtype: %d, size: %d, v: %f, w: %f, target_speed: %f, remain_dis: %f, error: %f %f", 
    //           running_type_, (int)trajectory.points.size(), speed, angle_speed, trajectory.points[0].target_speed, trajectory.points.back().distance,
    //           amr_state.lat_error, amr_state.heading_error);
}

void LQRPathFollowing::HeadFollowingController(amr_trajectory_msgs::amr_trajectory_msg &trajectory,
                                               AMRState &amr_state,
                                               float &speed, float &angle_speed,
                                               float &max_speed_limit, float &max_angle_speed_limit)
{
    float planning_speed = 0.0f;
    float lon_preveiwer_dis = g_lon_preveiwer_dis;

    if (lon_preveiwer_dis > trajectory.points.back().distance)
    {
        lon_preveiwer_dis = trajectory.points.back().distance;
    }

    for (auto p : trajectory.points)
    {
        if (p.distance > lon_preveiwer_dis)
        {
            planning_speed = p.target_speed;
            break;
        }
    }
    // ROS_INFO("Planning_speed: line: %d, v: %f", __LINE__, planning_speed);
    float pose_control_speed = fabs(planning_speed);
    double line_angle = com_.AngleConstraintToPI(tf::getYaw(trajectory.points[0].pose.orientation));
    float remaind_dis = TargetLongDistance(amr_state, trajectory.points.back(), line_angle);

    if (trajectory.points.back().distance < g_slow_down_distance)
    {
        if (pose_control_speed > g_slow_down_speed)
        {
            pose_control_speed = g_slow_down_speed;
        }
    }

    if (trajectory.points.back().distance < g_position_control_distance)
    {
        float speed_temp = g_position_control_kp * trajectory.points.back().distance;
        if (pose_control_speed > speed_temp)
        {
            pose_control_speed = speed_temp;
        }
    }

    if (pose_control_speed < g_slow_down_speed) // g_min_move_speed
    {
        pose_control_speed = g_slow_down_speed;
    }
    
    if (stop_distance_threshold_.GetState(remaind_dis) == common_lib::DoubleThreshold_e::DOEN_AREA)
    {
        pose_control_speed = 0;
        ptr_speed_smooth_->Reset();
        LINFO("LQR --> DONE");
    }
    else
    {
        if (fabs(planning_speed) < 0.0001)
        {
            // 需将最后一点的坐标转换到载体系下
            Eigen::Vector3d target = Eigen::Vector3d(trajectory.points.back().pose.position.x, trajectory.points.back().pose.position.y, 0.0f);
            Eigen::Vector3d cali_pose = tra_.PosTF(target, amr_state.amr_local_.pose);

            // LINFO("LQR --> LAST x: {:0.6f}, y: {:0.6f} dis: {:0.6f}", cali_pose(0), cali_pose(1), std::hypot(cali_pose(0), cali_pose(1)));
            pose_control_speed = g_slow_down_speed * 1.1f;
            if (cali_pose(0) > 0.0f)
            {
                planning_speed = 1.0;
            }
            else
            {
                planning_speed = -1.0;
            }
        }
    }
    // distance_dec_ = std::max(0.0, double(remaind_dis - g_stop_offset));
    distance_dec_ = remaind_dis;
    planning_speed = ((planning_speed > 0.0f) - (planning_speed < 0.0f)) * pose_control_speed;
    // ROS_INFO("planning_speed: %f", planning_speed);
    LineSpeedLimit(amr_state, planning_speed, max_speed_limit);
    speed = planning_speed;

    if (fabs(planning_speed) < FLT_EPSILON || fabs(amr_state.actuator.walking_speed) < g_slow_down_speed) // g_stop_adjust_walk_speed)
    {
        LINFO("LQR --> vehicle slow move and path following set angle speed to 0");
        angle_speed = 0;
    }
    else
    {
        for (auto point : trajectory.points)
        {
            if (0.000 < point.distance)
            {
                float loc_theta = tf::getYaw(point.loc_pose.orientation);
                if (speed < 0) // back
                {
                    loc_theta = com_.AngleConstraintToPI(loc_theta + M_PI);
                    point.loc_pose.position.y = -point.loc_pose.position.y;
                    planning_speed = -planning_speed;
                }

                Eigen::Matrix4f A;
                Eigen::Vector4f B;

                Eigen::Matrix4f Q;
                Eigen::Vector4f X;

                Q.setZero();

                if (fabs(point.curvature) > FLT_EPSILON)
                {
                    Q = LQR_Q_load_;
                }
                else
                {
                    Q = LQR_Q_unload_;
                }

                float R = 1;
                A.setZero();
                float dt = 0.05f;
                A(0, 0) = 1.0f;
                A(0, 1) = dt;
                A(1, 2) = planning_speed;
                A(2, 2) = 1.0f;
                A(2, 3) = dt;

                B.setZero();
                B[3] = -1;

                static float pe_ = point.loc_pose.position.y;
                static float pth_e = loc_theta;
                X[0] = point.loc_pose.position.y;
                X[1] = (point.loc_pose.position.y - pe_) / dt;
                X[2] = loc_theta;
                X[3] = (loc_theta - pth_e) / dt;
                pe_ = point.loc_pose.position.y;
                pth_e = loc_theta;

                Eigen::Matrix4f P = DLQR(A, B, Q, R);
                Eigen::RowVector4f K = 1.0 / (B.transpose() * P * B + R) * (B.transpose() * P * A);

                static float last_curv = point.curvature;
                angle_speed = -K * X + point.curvature * fabs(planning_speed) * g_curvatrue_control_weight;
                LINFO("lat: {:0.6f}, lat_dt: {:0.6f}, heading: {:0.6f}, heading_dt: {:0.6f}, {:0.6f}, curvature: {:0.6f}, cs: {:0.6f}, , ret: {:0.6f}", 
                      -K[0]*X[0], -K[1]*X[1], -K[2]*X[2], -K[3]*X[3], (float)(-K*X), point.curvature, point.curvature*fabs(planning_speed),  angle_speed);

                last_curv = point.curvature;

                angle_speed = com_.ValueConstraintMinToMax(angle_speed, -g_max_control_angle_speed, g_max_control_angle_speed);
                angle_speed = com_.ValueConstraintMinToMax(angle_speed, -fabs(max_angle_speed_limit), fabs(max_angle_speed_limit));
                amr_state.lat_error = point.loc_pose.position.y;
                amr_state.heading_error = loc_theta;

                if (angle_speed != 0)
                {
                    if (fabs(speed / angle_speed) < g_headfollowing_min_radius)
                    {
                        angle_speed = (fabs(speed) / g_headfollowing_min_radius) * angle_speed / fabs(angle_speed);
                        LWARNING("LQR --> Limit radius: v: {}, w: {}", speed, angle_speed);
                    }
                }
                break;
            }
        }
    }
    
    // ROS_INFO("Planning_speed: line: %d, v: %f, dis: %f", __LINE__, speed, remaind_dis);
}

float LQRPathFollowing::TargetLongDistance(AMRState &amr_state, amr_trajectory_msgs::point_info target_point, float line_theta)
{
    float dis;
    float heading;
    if (target_point.target_speed > 0)
    {
        heading = 0;
    }
    else
    {
        heading = M_PI;
    }
    float theta = tf::getYaw(amr_state.amr_local_.pose.orientation);
    float global_steer_angle = com_.AngleConstraintToPI(theta + heading);
    global_steer_angle = line_theta;
    dis = (target_point.pose.position.x - amr_state.amr_local_.pose.position.x) * cos(global_steer_angle) + (target_point.pose.position.y - amr_state.amr_local_.pose.position.y) * sin(global_steer_angle); // + 0.007;
    if (dis < 0)
    {
        return 0.0f;
    }
    else
    {
        return dis;
    }
}

void LQRPathFollowing::LineSpeedLimit(AMRState &amr_state, float &speed, float max_speed_limit)
{
    float target_speed = speed;
    if (fabs(target_speed) > max_speed_limit)
    {
        target_speed = ((target_speed>FLT_EPSILON) - (target_speed<-FLT_EPSILON)) * max_speed_limit;
    }

    // StartSpeed
    if (fabs(amr_state.actuator.walking_speed) < g_zero_speed)
    {
        if ( g_zero_speed < fabs(target_speed) && g_min_move_speed > fabs(target_speed))
        {
            int sign = target_speed > 0 ? 1: -1;
            target_speed = sign * g_min_move_speed;
        }
    }
    // ROS_INFO("LineSpeedLimit(): v: %f", amr_state.actuator.walking_speed);
    float walking_speed = (fabs(amr_state.actuator.walking_speed)<g_min_move_speed*0.6f)?0.0f:amr_state.actuator.walking_speed;
    float speed_max_acc = g_max_acc;
    float speed_max_dcc = g_max_dcc;

    speed = ptr_speed_smooth_->VelocityConfig(target_speed, walking_speed, speed_max_acc, speed_max_dcc);
    
    float calc_dec_velocity = sqrt(fabs(2 * speed_max_dcc * distance_dec_)); // v^2 = 2as

    // 计算当前位置的理论速度与实际速度差值
    float vel_error    = std::max(fabs(amr_state.actuator.walking_speed) - calc_dec_velocity, 0.0f);
    float vel_error_k  = g_dec_vel_error_k;
    float time_delay   = g_dec_response_delay;
    float pre_distance = distance_dec_ - fabs(amr_state.actuator.walking_speed) * (time_delay + vel_error / speed_max_dcc * vel_error_k);
    float dec_speed = ptr_speed_smooth_->VelocityConfigToEnd(pre_distance, walking_speed, speed_max_acc, speed_max_dcc);
    
    // ROS_INFO("Planning_speed: line: %d, v: %f, dec_speed: %f, distance: %f %f", __LINE__, speed, dec_speed, distance_dec_, pre_distance);
    if (fabs(dec_speed) < g_min_move_speed && distance_dec_ > 0)
    {
        dec_speed = g_min_move_speed;
    }
    
    if (fabs(dec_speed) < fabs(speed))
    {
        int sign = speed >= 0 ? 1 : -1;
        speed = ((speed>FLT_EPSILON) - (speed<-FLT_EPSILON)) * dec_speed;
        // ptr_speed_smooth_->SyncSpeed(speed, g_speed_plan_type);
    }
    // ROS_INFO("LineSpeedLimit: line: %d, v: %f %f", __LINE__, speed, target_speed);
}

void LQRPathFollowing::RotatingController(amr_trajectory_msgs::amr_trajectory_msg &trajectory, AMRState &amr_state, float &speed, float &angle_speed, float &max_angle_speed_limit)
{
    float delta_heading = 0.0f;
    float walk_heading = tf::getYaw(amr_state.amr_local_.pose.orientation);

    int cnt = 4;
    int index = cnt;
    if (cnt > trajectory.points.size() - 1)
    {
        index = trajectory.points.size() - 1;
    }
    
    float target_theta = trajectory.points[index].amr_angle;
    LINFO("index {}, path.points[index].angle {:0.6f}", index, tf::getYaw(trajectory.points[index].pose.orientation));
    
    int direction = com_.AngleConstraintToPI(target_theta - walk_heading) >= 0 ? 1 : -1;
    float max_speed = /*amr_state.actuator.load_status ? com_.Degree2Radian(g_max_rotating_angle_speed_degree_inload) : */g_max_rotating_angle_speed;
    angle_speed = direction * max_speed;
    
    angle_speed = com_.ValueConstraintMinToMax(angle_speed, -fabs(max_angle_speed_limit), fabs(max_angle_speed_limit));
    if (fabs(angle_speed) < g_min_rotating_angle_speed)
    {
        angle_speed = ((angle_speed > 0) - (angle_speed < 0)) * (g_min_rotating_angle_speed);
    }
    speed = 0;
    amr_state.lat_error = 0;
    amr_state.heading_error = com_.AngleConstraintToPI(target_theta - walk_heading); 

    // Angle Speed Smooth
    float min_angle_speed = (g_min_rotating_angle_speed);
    if (0.005 < fabs(angle_speed) && min_angle_speed > fabs(angle_speed) && fabs(amr_state.actuator.angle_speed) < 0.005f)
    {
        int sign = angle_speed > 0 ? 1: -1;
        angle_speed = sign * min_angle_speed;
    }

    float target_angle_speed = angle_speed;
    float angle_max_acc = g_max_angle_speed_acc;
    float angle_max_dcc = g_max_angle_speed_dcc;
    float angle_speed_state = (fabs(amr_state.actuator.angle_speed)<min_angle_speed*0.6f)?0.0f:amr_state.actuator.angle_speed;
    angle_speed = ptr_anglespeed_smooth_->VelocityConfig(target_angle_speed, 
                                                              angle_speed_state,
                                                              angle_max_acc,
                                                              angle_max_dcc);

    float angle_speed_tmp = ptr_anglespeed_smooth_->VelocityConfigToEnd(fabs(amr_state.heading_error), angle_speed_state, angle_max_acc, angle_max_dcc);
    // ROS_INFO("Target angle speed: %f %f, acc: %f, state: %f", angle_speed, angle_speed_tmp, angle_max_acc, angle_speed_state);
    angle_speed = direction * fabs(std::min(angle_speed, angle_speed_tmp));
    // LINFO("Rotate Controller");
    ROS_INFO("RotatingController() --> dir: %d, target : %f, now: %f, w: %f", direction, target_theta, walk_heading, angle_speed);
}

bool LQRPathFollowing::IsFollowingToRotating(amr_trajectory_msgs::amr_trajectory_msg &trajectory, AMRState &amr_state)
{
    int num = 3;
    int cnt = 0;
    float d_heading = 0;
    float walk_heading;// = tf::getYaw(amr_state.amr_local_.pose.orientation);

    float heading_error = 0.0f;

    if (trajectory.points.size() == 1)
    {
        heading_error = com_.Degree2Radian(1.0f);
    }
    else
    {
        heading_error = g_f2r_heading_error;
    }

    while (cnt < num && cnt < trajectory.points.size())
    {
        d_heading = 0.0f;
        walk_heading = tf::getYaw(amr_state.amr_local_.pose.orientation);
        float heading = trajectory.points[cnt].amr_angle;

        if (fabs(com_.AngleConstraintToPI(walk_heading - heading)) > heading_error)
        {
            LINFO("IsFollowingToRotating() ");
            return true;
        }
        cnt ++;
    }
    return false;
}

bool LQRPathFollowing::IsRotatingToFollowing(amr_trajectory_msgs::amr_trajectory_msg &trajectory, AMRState &amr_state)
{
    float d_heading = 0;
    float walk_heading = tf::getYaw(amr_state.amr_local_.pose.orientation);
    int cnt = 4;
    int index = cnt;
    if (cnt > trajectory.points.size() - 1)
    {
        index = trajectory.points.size() - 1;
    }

    if (fabs(com_.AngleConstraintToPI((walk_heading - trajectory.points[index].amr_angle))) < g_r2f_heading_error)
    {
        return true;
    }
    return false;
}

Eigen::Matrix4f LQRPathFollowing::DLQR(Eigen::Matrix4f A, Eigen::Vector4f B, Eigen::Matrix4f Q, float R)
{
    Eigen::Matrix4f X = Q;
    Eigen::Matrix4f Xn;
    int maxiter = 500; // g_eps_maxiter;
    float eps = 0.01f; // g_eps;
    int i = 0;

    auto A_inv = A.transpose();
    auto B_inv = B.transpose();

    for (i = 0; i < maxiter; i++)
    {
        // Xn = A.transpose() * X * A - A.transpose() * X * B / (R + B.transpose() * X * B) * B.transpose() * X * A + Q;
        Xn = A_inv * X * A - A_inv* X * B / (R + B_inv * X * B) * B_inv * X * A + Q;
        Eigen::Matrix4f error = Xn - X;
        if (error.cwiseAbs().maxCoeff() < eps)
        {
            return Xn;
        }
        X = Xn;
    }
    
    return Xn;
}